#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
import math
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel


def render(viewpoint_camera, 
           pc : GaussianModel, 
           pipe, 
           bg_color : torch.Tensor, 
           scaling_modifier = 1.0, 
           override_color = None,
           cur_iters=-1):
    """
    Render the scene. 
    
    Background tensor (bg_color) must be on GPU!
    """

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
    except:
        pass
    tem_opacity = torch.zeros_like(pc.get_xyz[:, :1], dtype=pc.get_xyz.dtype, device="cuda") + 0

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=0,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug)

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
    
    all_means3D = pc.get_xyz
    all_means2D = screenspace_points
    all_scales = pc.get_scaling
    all_rotations = pc.get_rotation

    visible_mask = rasterizer.markVisible(positions=all_means3D)
    colors, opacity, remap_opacity, view_opacity, mask = pc.feature_encode(viewpoint_camera, visible_mask, cur_iters)
    
    # Rasterize visible Gaussians to image, obtain their radii (on screen). 
    rendered_image, static_radii = rasterizer(means3D=all_means3D[visible_mask][mask],
                                        means2D=all_means2D[visible_mask][mask],
                                        shs=None,
                                        colors_precomp=colors[mask],
                                        opacities=opacity,
                                        scales=all_scales[visible_mask][mask],
                                        rotations=all_rotations[visible_mask][mask],
                                        cov3D_precomp=None)
    
    rendered_occlu, occlu_radii, = rasterizer(means3D=all_means3D[visible_mask][~mask],
                                        means2D=all_means2D[visible_mask][~mask],
                                        shs=None,
                                        colors_precomp=colors[~mask],
                                        opacities=remap_opacity,
                                        scales=all_scales[visible_mask][~mask],
                                        rotations=all_rotations[visible_mask][~mask],
                                        cov3D_precomp=None)
        
    view_radii = torch.zeros_like(mask, dtype=static_radii.dtype, device=mask.device)
    view_radii[mask] = static_radii
    view_radii[~mask] = occlu_radii
    tem_radii = torch.zeros_like(visible_mask, dtype=static_radii.dtype, device=visible_mask.device)
    tem_radii[visible_mask] = view_radii
    visibility_filter = torch.logical_and(visible_mask, tem_radii>0)
    tem_opacity[visible_mask] = view_opacity 

    return {
        "render": rendered_image + rendered_occlu,
        'render_stastic': rendered_image,
        'render_occlu': rendered_occlu,
        "viewspace_points": screenspace_points,
        "visibility_filter" : visibility_filter,
        "view_opacity": tem_opacity,
        "radii": tem_radii,
        'static_opacity': all_scales[visible_mask],
        'occlu_opacity': remap_opacity
    }
